------------------------------------------------------------------------------
-- Copyright 2023, Tim Whiting, Fredrik Wieczerkowski
--
-- This is free software; you can redistribute it and/or modify it under the
-- terms of the Apache License, Version 2.0. A copy of the License can be
-- found in the LICENSE file at the root of this distribution.
-----------------------------------------------------------------------------

-----------------------------------------------------------------------------
-- The language server's monad that holds state (e.g. loaded/compiled modules)
-----------------------------------------------------------------------------
{-# LANGUAGE OverloadedStrings #-}

module LanguageServer.Monad
  ( LSState (..),
    InlayHintOptions(..),
    SignatureContext(..), clearSignatureContext, updateSignatureContext, getSignatureContext,
    Colors(..), updateColorScheme, getColorScheme, toJSON, fromJSON,
    defaultLSState,
    newLSStateVar,
    LSM,
    getTerminal, getFlags, getHtmlPrinter,
    getLSState,modifyLSState,
    updateConfig,
    getInlayHintOptions,
    runLSM,
    getProgress, setProgress, maybeContents,

    liftBuild, liftBuildWith,
    lookupModuleName, lookupRangeMap, lookupProgram, lookupLexemes,
    lookupDefinitions, lookupVisibleDefinitions, Definitions(..),
    lookupModulePaths,
    getPrettyEnv, getPrettyEnvFor, prettyMarkdown,
    emitInfo, emitNotification, getVirtualFileVersion

  )
where

import Debug.Trace(trace)
import GHC.Conc (atomically)
import GHC.Base (Alternative (..), when)
import Platform.Var (newVar, takeVar)
import Platform.Filetime (FileTime)
import Control.Concurrent.MVar (MVar, modifyMVar_, newMVar, readMVar)
import Control.Monad.Reader (ReaderT, ask, runReaderT)
import Control.Monad.Trans (lift, liftIO)
import Control.Concurrent.Chan (readChan)
import Control.Concurrent.STM (newTVarIO, TVar)
import qualified Data.Set as Set
import Control.Concurrent.STM.TChan
import qualified Data.Aeson as A
import qualified Data.ByteString as D
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Data.Maybe (fromMaybe, isJust, isNothing)
import Data.List (find)
import Data.Aeson.Types
import Language.LSP.Server (LanguageContextEnv, LspT, runLspT, sendNotification, Handlers)
import qualified Language.LSP.Protocol.Types as J
import qualified Language.LSP.Protocol.Message as J
import qualified Language.LSP.Server as J

import Lib.PPrint hiding (empty)
import qualified Lib.PPrint as PP
import Lib.Printer
import Common.ColorScheme ( colorSource, ColorScheme, darkColorScheme, lightColorScheme )
import Common.Name (nameNil, Name, readQualifiedName, ModuleName)
import Common.File ( realPath, normalize, getCwd, realPath, normalize, getCurrentTime )
import Common.Error

import Syntax.Syntax( UserProgram )
import Syntax.RangeMap( RangeMap )
import Syntax.Lexeme( Lexeme )
import Kind.ImportMap (importsEmpty)
import qualified Type.Pretty as TP
import Compile.Options (Flags (..), prettyEnvFromFlags, verbose, Terminal(..))
import Compile.BuildContext
import LanguageServer.Conversions ({-toLspUri,-} fromLspUri)

import Data.Map.Strict(Map)
import Data.ByteString (ByteString)
import GHC.IO.Encoding (BufferCodec(getState))
import Numeric

-- The language server's state, e.g. holding loaded/compiled modules.
data LSState = LSState {
  buildContext      :: !BuildContext,

  messages          :: !(TChan (String, J.MessageType)),
  progress          :: !(TChan (Double, String)),
  flags             :: !Flags,
  terminal          :: !Terminal,
  progressReport    :: !(Maybe (J.ProgressAmount -> LSM ())),
  htmlPrinter       :: !(Doc -> IO T.Text),

  pendingRequests   :: !(TVar (Set.Set J.SomeLspId)),
  cancelledRequests :: !(TVar (Set.Set J.SomeLspId)),
  documentVersions  :: !(TVar (M.Map J.NormalizedUri J.Int32)),
  documentInfos     :: !(M.Map J.NormalizedUri (D.ByteString, FileTime, J.Int32)),
  signatureContext  :: !(Maybe SignatureContext),
  config            :: !Config
}

data SignatureContext = SignatureContext {
  sigFunctionName :: !Name
}

instance FromJSON SignatureContext where
  parseJSON (A.Object v) = SignatureContext <$> readQualifiedName <$> v .: "sigFunctionName"
  parseJSON _ = return $ SignatureContext nameNil

instance ToJSON SignatureContext where
  toJSON (SignatureContext name) = object ["sigFunctionName" .= (show name)]

-- The monad holding (thread-safe) state used by the language server.
type LSM = LspT () (ReaderT (MVar LSState) IO)

-- Runs the language server's state monad.
runLSM :: LSM a -> MVar LSState -> LanguageContextEnv () -> IO a
runLSM lsm stVar cfg = runReaderT (runLspT cfg lsm) stVar

newLSStateVar :: Flags -> IO (MVar LSState)
newLSStateVar flags = defaultLSState flags >>= newMVar

-- Fetches the language server's state inside the LSM monad
getLSState :: LSM LSState
getLSState = do
  stVar <- lift ask
  liftIO $ readMVar stVar

-- Updates the language server's state inside the LSM monad
modifyLSState :: (LSState -> LSState) -> LSM ()
modifyLSState f = do
  stVar <- lift ask
  liftIO $ modifyMVar_ stVar (return . f)

defaultLSState :: Flags -> IO LSState
defaultLSState flags = do
  msgChan <- atomically newTChan :: IO (TChan (String, J.MessageType))
  progressChan <- atomically newTChan :: IO (TChan (Double, String))
  pendingRequests <- newTVarIO Set.empty
  cancelledRequests <- newTVarIO Set.empty
  fileVersions <- newTVarIO M.empty
  let withNewPrinter f = do
        ansiConsole <- newVar ansiDefault
        stringVar <- newVar ""
        let p = AnsiString ansiConsole stringVar
        tp <- (f . PAnsiString) p
        ansiString <- takeVar stringVar
        atomically $ writeTChan msgChan (ansiString, tp)
  let withNewProgressPrinter percent mbdoc = do
        let perc = 100 * percent
        atomically $ writeTChan progressChan
          (perc, show (fill 7 (text ((showFFloat (Just 2) perc) "%")) <.>
                       case mbdoc of
                         Just doc -> space <+> doc
                         Nothing  -> space)
          )
  cwd <- getCwd
  let cscheme = colorScheme flags
      prettyEnv flags ctx imports = (prettyEnvFromFlags flags){ TP.context = ctx, TP.importsMap = imports }
      term = Terminal (\err -> withNewPrinter $ \p -> do putErrorMessage p cwd (showSpan flags) cscheme err; return J.MessageType_Error)
                (if verbose flags > 1 then (\msg -> withNewPrinter $ \p -> do withColor p (colorSource cscheme) (writeLn p msg); return J.MessageType_Info)
                                         else (\_ -> return ()))
                 (\(percent, mbdoc) -> withNewProgressPrinter percent mbdoc)
                 (if verbose flags > 0 then (\msg ->
                    withNewPrinter $ \p -> do
                      writePrettyLn p msg
                      return J.MessageType_Info
                    )
                  else (\_ -> return ()))
                 (\msg -> withNewPrinter $ \p -> do
                    writePrettyLn p msg
                    return J.MessageType_Info
                )
  return LSState {
    buildContext = buildcEmpty flags,
    messages = msgChan, progress=progressChan, pendingRequests=pendingRequests, cancelledRequests=cancelledRequests,
    terminal = term, htmlPrinter = htmlTextColorPrinter, flags = flags,
    documentInfos = M.empty, documentVersions = fileVersions,
    signatureContext = Nothing, progressReport = Nothing,
    config = Config{
      langServerOpts = LanguageServerOptions{
        inlayHintOpts=InlayHintOptions{
                         showImplicitArguments=True,
                         showInferredTypes=True,
                         showFullQualifiers=True
        }
      }
    }
  }

-- Prints a message to html spans
htmlTextColorPrinter :: Doc -> IO T.Text
htmlTextColorPrinter doc
  = do
    stringVar <- newVar (T.pack "")
    let printer = PHtmlText (HtmlTextPrinter stringVar)
    -- writePretty printer doc
    writePrettyW printer 1000 doc
    takeVar stringVar

putScheme p env tp
  = writePrettyLn p (TP.ppScheme env tp)

putErrorMessage p cwd endToo cscheme err
  = writePrettyLn p (ppErrorMessage cwd endToo cscheme err)

data KokaConfig = KokaConfig {
  kokaCfg :: Config
} deriving Show

instance FromJSON KokaConfig where
  parseJSON (A.Object v) = KokaConfig <$> v .: "koka"
  parseJSON _ = empty

data Config = Config {
  langServerOpts :: LanguageServerOptions
} deriving Show

instance FromJSON Config where
  parseJSON (A.Object v) = Config <$> v .: "languageServer"
  parseJSON _ = empty

data LanguageServerOptions = LanguageServerOptions {
  inlayHintOpts :: InlayHintOptions
} deriving Show

instance FromJSON LanguageServerOptions where
  parseJSON (A.Object v) = LanguageServerOptions <$> v .: "inlayHints"
  parseJSON _ = empty

data InlayHintOptions  = InlayHintOptions {
  showImplicitArguments :: Bool,
  showInferredTypes :: Bool,
  showFullQualifiers :: Bool
} deriving Show

instance FromJSON InlayHintOptions where
  parseJSON (A.Object v) = InlayHintOptions <$> v .: "showImplicitArguments" <*> v .: "showInferredTypes" <*> v .: "showFullQualifiers"
  parseJSON _ = empty

data Colors = Colors {
  mode :: String
} deriving Show

instance FromJSON Colors where
  parseJSON (A.Object v) = Colors <$> v .: "mode"
  parseJSON _ = empty


setProgress :: Maybe (J.ProgressAmount -> LSM ()) -> LSM ()
setProgress report = do
  modifyLSState $ \s -> s {progressReport = report}

getProgress :: LSM (Maybe (J.ProgressAmount -> LSM ()))
getProgress = progressReport <$> getLSState

updateConfig :: A.Value -> LSM ()
updateConfig cfg =
  case fromJSON cfg of
    A.Success cfg -> do
      -- trace ("Updating config to " ++ show cfg) $ return ()
      modifyLSState $ \s -> s{config=kokaCfg cfg}
    c -> 
      trace ("Failed to parse config " ++ show cfg ++ " error: " ++ show c) $
      return ()

updateColorScheme :: Colors -> LSM ()
updateColorScheme colors =
  modifyLSState $ \s ->
    trace ("Updating color scheme to " ++ mode colors) $
    case mode colors of
      "dark" -> s{flags=(flags s){colorScheme=darkColorScheme}}
      _ -> s{flags=(flags s){colorScheme=lightColorScheme}}

updateSignatureContext :: SignatureContext -> LSM ()
updateSignatureContext context =
  modifyLSState $ \s -> s{signatureContext=Just context}

clearSignatureContext :: LSM ()
clearSignatureContext =
  modifyLSState $ \s -> s{signatureContext=Nothing}

getSignatureContext :: LSM (Maybe SignatureContext)
getSignatureContext = signatureContext <$> getLSState

getInlayHintOptions :: LSM InlayHintOptions
getInlayHintOptions = inlayHintOpts . langServerOpts . config <$> getLSState

getVirtualFileVersion :: J.NormalizedUri -> LSM (Maybe J.Int32)
getVirtualFileVersion uri
  = do ls <- getLSState
       case M.lookup uri (documentInfos ls) of
         Just (_,_,ver) -> return (Just ver)
         _              -> return Nothing


-- Fetches the terminal used for printing messages
getTerminal :: LSM Terminal
getTerminal = terminal <$> getLSState

-- Fetches the loaded state holding compiled modules
getFlags :: LSM Flags
getFlags = flags <$> getLSState

-- Fetches the html printer used for printing markdown compatible text
getHtmlPrinter :: LSM (Doc -> IO T.Text)
getHtmlPrinter = htmlPrinter <$> getLSState

-- Fetches the color scheme used for coloring markdown compatible text
getColorScheme :: LSM ColorScheme
getColorScheme = colorScheme <$> getFlags



-- Retreives a file from the virtual file system, returning the contents and the last modified time
maybeContents :: Map J.NormalizedUri (ByteString, FileTime, J.Int32) -> FilePath -> Maybe (ByteString, FileTime)
maybeContents vfs path = do
  -- trace ("Maybe contents " ++ show uri ++ " " ++ show (M.keys vfs)) $ return ()
  let uri = J.toNormalizedUri $ J.filePathToUri path
  (text, ftime, vers) <- M.lookup uri vfs
  return (text, ftime)


-- Run a build with optionally temporarily changed flags (this restores the original build context afterwards)
-- (when a build context is validated, it checks itself against the current flags and rebuilds accordingly)
liftBuildWith :: Maybe Flags -> (BuildContext -> Build (BuildContext,a)) -> LSM (Either Errors (a,Errors))
liftBuildWith mbFlags action
  = do ls <- getLSState
       let vfs  = VFS (\fpath -> maybeContents (documentInfos ls) fpath)
           flgs = case mbFlags of
                    Nothing    -> flags ls
                    Just flags -> flags
       res <- seq flgs $ seq VFS $ liftIO $ runBuild (terminal ls) flgs $ withVFS vfs $ action (buildContext ls)
       case res of
         Left errs               -> return (Left errs)
         Right ((buildc,x),errs) -> do when (isNothing mbFlags) $
                                          modifyLSState (\ls -> ls{ buildContext = buildc })
                                       return (Right (x,errs))


-- Run a build monad with current terminal, flags, and virtual file system
liftBuild :: (BuildContext -> Build (BuildContext,a)) -> LSM (Either Errors (a,Errors))
liftBuild action
  = liftBuildWith Nothing action

getBuildContext :: LSM BuildContext
getBuildContext
  = do ls <- getLSState
       return (buildContext ls)

-- Module name from URI
lookupModuleName :: J.NormalizedUri -> LSM (Maybe (FilePath,ModuleName))
lookupModuleName uri
  = do buildc <- getBuildContext
       mbfpath  <- liftIO $ fromLspUri uri
       return $ do fpath   <- mbfpath -- maybe monad
                   modname <- buildcLookupModuleName fpath buildc
                   return (fpath,modname)

-- Lexemes from module name
-- available even if the source cannot be parsed
lookupLexemes :: ModuleName -> LSM (Maybe [Lexeme])
lookupLexemes mname
  = do buildc <- getBuildContext
       return (buildcGetLexemes mname buildc)

-- RangeMap from module name
lookupRangeMap :: ModuleName -> LSM (Maybe (RangeMap,[Lexeme]))
lookupRangeMap mname
  = do buildc <- getBuildContext
       return (buildcGetRangeMap mname buildc)

-- Program from module name
lookupProgram :: ModuleName -> LSM (Maybe UserProgram)
lookupProgram mname
  = do buildc <- getBuildContext
       return (buildcLookupProgram mname buildc)


-- Pretty environment
getPrettyEnv :: LSM TP.Env
getPrettyEnv
  = do flags <- getFlags
       return (prettyEnvFromFlags flags)

-- Pretty environment
getPrettyEnvFor :: ModuleName -> LSM TP.Env
getPrettyEnvFor modname
  = do flags <- getFlags
       return (prettyEnvFromFlags flags){ TP.context = modname }


-- Format as markdown
prettyMarkdown :: Doc -> LSM T.Text
prettyMarkdown doc
  = do htmlPrinter <- getHtmlPrinter
       liftIO (htmlPrinter doc)

-- Sent text to the terminal info.
emitInfo :: (TP.Env -> Doc) -> LSM ()
emitInfo mkDoc
  = do penv <- getPrettyEnv
       term <- getTerminal
       liftIO $ termInfo term (mkDoc penv)

-- Emit an error notification
emitNotification :: (TP.Env -> Doc) -> LSM ()
emitNotification mkDoc
  = do penv     <- getPrettyEnv
       markdown <- prettyMarkdown (mkDoc penv)
       sendNotification J.SMethod_WindowLogMessage $ J.LogMessageParams J.MessageType_Error markdown

-- Return definitions (gamma etc) for a set of modules
lookupDefinitions :: [ModuleName] -> LSM Definitions
lookupDefinitions modnames
  = do buildc <- getBuildContext
       return (buildcGetDefinitions modnames buildc)

-- Return definitions (gamma etc) for a set of modules including the imports.
lookupVisibleDefinitions :: [ModuleName] -> LSM Definitions
lookupVisibleDefinitions modnames
  = do buildc <- getBuildContext
       return (buildcGetVisibleDefinitions modnames buildc)

-- Return all loaded module names and associated file paths.
lookupModulePaths :: LSM [(ModuleName,FilePath)]
lookupModulePaths
  = do buildc <- getBuildContext
       return (buildcModulePaths buildc)
